﻿#pragma once

#include <xy/alg/bounding_box.h>
#include <xy/alg/coeff_matrix.h>
#include <xy/alg/iterable.h>
#include <xy/alg/projection.h>

namespace xy
{

namespace alg
{

template <std::size_t N> struct pia_curve : iterable<std::shared_ptr<geom::bspline_curve<N>>>
{
    using vec_t = geom::vec<N, double>; // P

    std::size_t degree{3};
    std::vector<vec_t> data;
    std::vector<double> params;

  protected:
    double m_w;
    std::shared_ptr<geom::bspline_curve<N>> m_curve;

    void init() override
    {
        if (params.empty())
            params = geom::chord_param(data);

        auto knots = geom::average_knots(params, degree);
        std::vector<vec_t> poles(data.size());
        for (std::size_t i = 0; i < data.size(); i++)
            poles[i] = data[i];

        // 计算全正基配置矩阵的最小特征值
        math::matxd A(data.size(), data.size());
        A.set_zero();
        for (std::size_t i = 0; i < A.rows(); i++)
        {
            auto span = geom::find_span(knots.data(), knots.size(), degree, params[i]);
            auto basis = geom::basis_funcs(knots.data(), degree, span, params[i]);
            for (std::size_t j = 0; j < degree + 1; j++)
                A(i, span - degree + j) += basis[j];
        }

        // 最佳权重 w = 2 / (1 + lambda(A)), 其中 lambda(A) \in (0,1) 是 A 的最小特征值
        for (std::size_t i = 0; i < A.rows(); i++)
            A(i, i) -= 1;
        m_w = 2 / (2 - math::eigenvalue(A));

        m_curve = std::make_shared<geom::bspline_curve<N>>(std::move(poles), std::move(knots), degree);
    }

    std::shared_ptr<geom::bspline_curve<N>> next() override
    {
        std::vector<vec_t> nPoles(data.size());
        for (std::size_t i = 0; i < nPoles.size(); i++)
            nPoles[i] = m_curve->pole(i) + (data[i] - m_curve->d0(params[i])) * m_w;

        m_curve = std::make_shared<geom::bspline_curve<N>>(std::move(nPoles), m_curve->move_knots(), degree);
        return m_curve;
    }

    std::shared_ptr<geom::bspline_curve<N>> next_preserve() override
    {
        std::vector<vec_t> nPoles(data.size());
        for (std::size_t i = 0; i < nPoles.size(); i++)
            nPoles[i] = m_curve->pole(i) + (data[i] - m_curve->d0(params[i])) * m_w;

        auto knots = m_curve->knots();
        m_curve = std::make_shared<geom::bspline_curve<N>>(std::move(nPoles), std::move(knots), degree);
        return m_curve;
    }

    bool should_continue() override
    {
        for (std::size_t i = 0; i < data.size(); i++)
        {
            auto t = projection_point<N>::nearest(m_curve, data[i], params[i], math::numeric::half);
            if ((m_curve->d0(t) - data[i]).modulus() > math::numeric::half)
                return true;
        }
        return false;
    }
};

template <std::size_t N> struct lspia_curve : iterable<std::shared_ptr<geom::bspline_curve<N>>>
{
    using vec_t = geom::vec<N, double>; // P

    std::size_t degree{3};
    std::size_t poleNum;
    std::vector<vec_t> data;
    std::vector<double> params;
    double error{1}; // 记录每次迭代的最大误差

  protected:
    std::shared_ptr<geom::bspline_curve<N>> m_curve;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = geom::chord_param(data);

        // 寻找最大曲率特征点
        std::vector<std::tuple<int, double>> curvatures;
        for (std::size_t i = 1; i < data.size() - 1; i++)
        {
            auto p1 = data[i - 1];
            auto p2 = data[i];
            auto p3 = data[i + 1];
            auto v1 = p2 - p1;
            auto v2 = p1 - p2 * 2 + p3;
            curvatures.emplace_back(i, cross(v1, v2).modulus() / std::pow(v1.modulus(), 3));
        }
        std::sort(curvatures.begin(), curvatures.end(),
                  [](auto &a, auto &b) { return std::get<1>(a) > std::get<1>(b); });
        std::sort(curvatures.begin(), curvatures.begin() + poleNum - 2,
                  [](auto &a, auto &b) { return std::get<0>(a) < std::get<0>(b); });

        std::vector<vec_t> poles(poleNum);
        std::vector<double> params(poleNum);
        poles[0] = data.front();
        params[0] = params.front();
        for (std::size_t i = 1; i < poleNum - 1; i++)
        {
            auto j = std::get<0>(curvatures[i - 1]);
            poles[i] = data[j];
            params[i] = params[j];
        }
        poles[poleNum - 1] = data.back();
        params[poleNum - 1] = params.back();

        // 使用这些控制点对应的参数值来构造节点向量
        auto knots = geom::average_knots(params, degree);
        m_curve = std::make_shared<geom::bspline_curve<N>>(std::move(poles), std::move(knots), degree);
    }

    /**
     * @brief 迭代计算下一个曲线
     *
     * @return std::shared_ptr<geom::bspline_curve<N>>
     */
    std::shared_ptr<geom::bspline_curve<N>> next() override
    {
        auto m = data.size() - 1;
        auto n = poleNum - 1;

        std::vector<double> basisSum(n + 1, 0); // 记录每个控制点对应的那组基函数和
        std::vector<vec_t> deltaSum(n + 1);     // 记录每个控制点对应的那组偏差向量加权和

        // 对每个数据点计算增量
        auto knots = m_curve->knots();
        for (std::size_t i = 0; i < m + 1; i++)
        {
            auto delta = data[i] - m_curve->d0(params[i]);

            // 计算此参数对应的非零基函数值
            auto span = geom::find_span(knots.data(), knots.size(), degree, params[i]);
            auto basis = geom::basis_funcs(knots.data(), degree, span, params[i]);

            // 累计参数
            for (std::size_t k = 0; k < degree + 1; k++)
            {
                basisSum[span - degree + k] += basis[k];
                deltaSum[span - degree + k] += delta * basis[k];
            }
        }

        // 平均基函数作为权重
        for (std::size_t i = 0; i < n + 1; i++)
            m_curve->set_pole(m_curve->pole(i) + deltaSum[i] / basisSum[i], i);
        return m_curve;
    }

    /**
     * @brief 根据逼近精度判断是否应该继续迭代
     *
     * @return true
     * @return false
     */
    bool should_continue() override
    {
        double err = 0;
        std::size_t max = 0;
        for (std::size_t i = 0; i < data.size(); i++)
        {
            auto t = projection_point<N>::nearest(m_curve, data[i], params[i], math::numeric::half);
            auto d = (m_curve->d0(t) - data[i]).modulus();
            if (d > err)
            {
                err = d;
                max = i;
            }
        }
        // 限制不取右端点
        max = std::min(max, data.size() - 2);

        if (err > math::numeric::half && std::abs(error - err) > math::numeric::half)
        {
            error = err;

            // 在最大误差点所在区间插入节点，不能让控制点数超过数据点数
            if (m_curve->pole_num() < data.size())
            {
                // Q \in [u_{span}, u_{span+1}]
                auto span =
                    geom::find_span(m_curve->knots().data(), m_curve->knots().size(), m_curve->degree(), params[max]);
                auto upper =
                    std::lower_bound(params.begin() + max, params.end(), m_curve->knots()[span + 1]) - params.begin();

                // 找到第一个大于等于 u_{span+1} 的参数位置
                std::vector<double> nknots;
                for (std::size_t i = max + 1; i < upper + 1; i++)
                {
                    double t = (params[i] + params[i - 1]) / 2;
                    m_curve->insert_knot(t);
                    poleNum++;

                    if (poleNum == data.size())
                        break;
                }
            }
            return true;
        }
        return false;
    }
};

template <std::size_t N> struct fairing_pia_curve : lspia_curve<N>
{
    using vec_t = geom::vec<N, double>; // P
    using lspia_curve<N>::degree;
    using lspia_curve<N>::poleNum;
    using lspia_curve<N>::data;
    using lspia_curve<N>::params;
    using lspia_curve<N>::error;

    std::vector<double> wStretch; // 每个控制点的拉伸权重
    std::vector<double> wBending; // 每个控制点的光顺权重

  protected:
    using lspia_curve<N>::m_curve;

    // 迭代格式 P' = A * P + B
    Eigen::SparseMatrix<double> m_A;
    Eigen::MatrixXd m_B;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        lspia_curve<N>::init();

        // 计算控制点的导数和二阶导数
        auto B0 = coeff_matrix::curve_least_square(params, m_curve->knots(), degree);
        auto [B1, B2] = coeff_matrix::curve_stretch_bending(m_curve->knots(), degree);
        auto [A, B] = coeff_matrices(B0, B1, B2, wStretch, wBending, data);

        m_A = std::move(A);
        m_B = std::move(B);
    }

    /**
     * @brief 迭代计算下一个曲线
     *
     * @return std::shared_ptr<geom::bspline_curve<N>>
     */
    std::shared_ptr<geom::bspline_curve<N>> next() override
    {
        auto n = poleNum - 1;

        // 填充控制点矩阵
        Eigen::MatrixXd P(n + 1, N);
        for (std::size_t i = 0; i < n + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                P(i, j) = m_curve->pole(i)[j];

        P = m_A * P + m_B;
        for (std::size_t i = 0; i < n + 1; i++)
        {
            vec_t p;
            for (std::size_t j = 0; j < N; j++)
                p[j] = P(i, j);
            m_curve->set_pole(p, i);
        }

        return m_curve;
    }

    /**
     * @brief 根据插值精度判断是否应该继续迭代
     *
     * @return true
     * @return false
     */
    bool should_continue() override
    {
        return true;
    }

  public:
    /**
     * @brief 根据控制点的拟合、拉伸和光顺矩阵计算系数矩阵
     *
     * @param[in] B0
     * @param[in] B1
     * @param[in] B2
     * @param[in] wStretch
     * @param[in] wBending
     * @param[in] data
     * @return std::tuple<Eigen::SparseMatrix<double>, Eigen::MatrixXd>
     */
    static std::tuple<Eigen::SparseMatrix<double>, Eigen::MatrixXd> coeff_matrices(
        const Eigen::SparseMatrix<double> &B0, const Eigen::SparseMatrix<double> &B1,
        const Eigen::SparseMatrix<double> &B2, const std::vector<double> &wStretch, const std::vector<double> &wBending,
        const std::vector<vec_t> &data)
    {
        auto n = wBending.size() - 1;
        std::vector<Eigen::Triplet<double>> tripletsO1, tripletsO2;
        for (int i = 0; i < n + 1; i++)
        {
            tripletsO1.push_back({i, i, wStretch[i]});
            tripletsO2.push_back({i, i, wBending[i]});
        }

        Eigen::SparseMatrix<double> O1(n + 1, n + 1), O2(n + 1, n + 1);
        O1.setFromTriplets(tripletsO1.begin(), tripletsO1.end());
        O2.setFromTriplets(tripletsO2.begin(), tripletsO2.end());

        // 计算系数矩阵
        Eigen::SparseMatrix<double> I(n + 1, n + 1);
        I.setIdentity();

        Eigen::SparseMatrix<double> A = (I - O1 - O2) * B0.transpose() * B0 + O1 * B1 + O2 * B2;

        // 构造权重 Lambda
        std::vector<Eigen::Triplet<double>> tripletsLam;

        // 计算行元素绝对值的和
        std::vector<double> rowSums(A.rows(), 0);
        for (int k = 0; k < A.outerSize(); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A, k); it; ++it)
                rowSums[it.row()] += std::abs(it.value());
        }
        // 这里必须用 int，不然类型不对
        for (int i = 0; i < rowSums.size(); i++)
            tripletsLam.push_back({i, i, 1 / rowSums[i]});

        Eigen::SparseMatrix<double> Lam(n + 1, n + 1);
        Lam.setFromTriplets(tripletsLam.begin(), tripletsLam.end());

        A = I - Lam * A;
        A.makeCompressed();

        // 填充点云矩阵
        Eigen::MatrixXd Q(data.size(), N);
        for (std::size_t k = 0; k < data.size(); k++)
            for (std::size_t j = 0; j < N; j++)
                Q(k, j) = data[k][j];

        // 最终的迭代矩阵
        return std::make_tuple(A, Lam * (I - O1 - O2) * B0.transpose() * Q);
    }
};

template <std::size_t N> struct pia_surface : iterable<std::shared_ptr<geom::bspline_surface<N>>>
{
    using vec_t = geom::vec<N, double>; // P

    std::size_t uDegree{3}, vDegree{3};
    std::vector<std::vector<vec_t>> data;
    std::vector<std::vector<geom::vec2>> params;

  protected:
    std::shared_ptr<geom::bspline_surface<N>> m_surface;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = geom::chord_param(data);

        std::vector<double> uParams(data.size(), 0);
        std::vector<double> vParams(data[0].size(), 0);

        // 平均参数值得到 u,v 方向的参数序列
        for (std::size_t i = 0; i < data.size(); i++)
        {
            for (std::size_t j = 0; j < data[0].size(); j++)
                uParams[i] += data[i][j].x();
            uParams[i] /= data[i].size();
        }

        for (std::size_t j = 0; j < data[0].size(); j++)
        {
            for (std::size_t i = 0; i < data.size(); i++)
                vParams[j] += data[i][j].y();
            vParams[j] /= data.size();
        }

        auto uKnots = geom::average_knots(uParams, uDegree);
        auto vKnots = geom::average_knots(vParams, vDegree);

        std::vector<vec_t> poles(data.size() * data[0].size());
        for (std::size_t i = 0; i < data.size(); i++)
            for (std::size_t j = 0; j < data[0].size(); j++)
                poles[i * data[0].size() + j] = data[i][j];

        m_surface = std::make_shared<geom::bspline_surface<N>>(std::move(poles), std::move(uKnots), std::move(vKnots),
                                                               uDegree, vDegree);
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<geom::bspline_surface<N>>
     */
    std::shared_ptr<geom::bspline_surface<N>> next() override
    {
        std::vector<vec_t> nPoles(data.size() * data[0].size());
        for (std::size_t i = 0; i < data.size(); i++)
        {
            for (std::size_t j = 0; j < data[0].size(); j++)
            {
                auto np = m_surface->pole(i, j) + data[i][j] - m_surface->d0(params[i][j]);
                nPoles[i * data[0].size() + j] = np;
            }
        }

        m_surface = std::make_shared<geom::bspline_surface<N>>(std::move(nPoles), m_surface->move_knots(0),
                                                               m_surface->move_knots(1), uDegree, vDegree);
        return m_surface;
    }
};

template <std::size_t N> struct lspia_surface : iterable<std::shared_ptr<geom::bspline_surface<N>>>
{
    using vec_t = geom::vec<N, double>; // P

    std::size_t uDegree{3}, vDegree{3};
    std::size_t uPoleNum, vPoleNum;
    std::vector<vec_t> data;
    std::vector<geom::vec2> params;
    std::shared_ptr<geom::bspline_surface<N>> surface;

  protected:
    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = bounding_box<N, double>::projection(data);

        if (surface == nullptr)
        {
            auto uKnots = geom::equally_knots(uPoleNum, uDegree);
            auto vKnots = geom::equally_knots(vPoleNum, vDegree);

            std::vector<vec_t> poles(uPoleNum * vPoleNum);
            surface = std::make_shared<geom::bspline_surface<N>>(std::move(poles), std::move(uKnots), std::move(vKnots),
                                                                 uDegree, vDegree);
        }
        else
        {
            uDegree = surface->degree(0);
            vDegree = surface->degree(1);
            uPoleNum = surface->pole_num(0);
            vPoleNum = surface->pole_num(1);
        }
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<geom::bspline_surface<N>>
     */
    std::shared_ptr<geom::bspline_surface<N>> next() override
    {
        auto m = data.size() - 1;
        auto n1 = uPoleNum - 1;
        auto n2 = vPoleNum - 1;

        std::vector<std::vector<double>> basisSum(n1 + 1,
                                                  std::vector<double>(n2 + 1, 0)); // 记录每个控制点对应的那组基函数和
        std::vector<std::vector<vec_t>> deltaSum(n1 + 1,
                                                 std::vector<vec_t>(n2 + 1)); // 记录每个控制点对应的那组偏差向量加权和

        // 对每个数据点计算增量
        auto uKnots = surface->knots(0);
        auto vKnots = surface->knots(1);
        for (std::size_t i = 0; i < m + 1; i++)
        {
            auto delta = data[i] - surface->d0(params[i]);

            // 计算此参数对应的非零基函数值
            auto uSpan = geom::find_span(uKnots.data(), uKnots.size(), uDegree, params[i].x());
            auto vSpan = geom::find_span(vKnots.data(), vKnots.size(), vDegree, params[i].y());
            auto uBasis = geom::basis_funcs(uKnots.data(), uDegree, uSpan, params[i].x());
            auto vBasis = geom::basis_funcs(vKnots.data(), vDegree, vSpan, params[i].y());

            // 累计参数
            for (std::size_t j = 0; j < uDegree + 1; j++)
            {
                auto row = uSpan - uDegree + j;
                for (std::size_t k = 0; k < vDegree + 1; k++)
                {
                    auto col = vSpan - vDegree + k;
                    basisSum[row][col] += uBasis[j] * vBasis[k];
                    deltaSum[row][col] += delta * uBasis[j] * vBasis[k];
                }
            }
        }

        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto pole = surface->pole(i, j);
                auto delta = deltaSum[i][j] / basisSum[i][j];
                surface->set_pole(pole + delta, i, j);
            }
        }
        return surface;
    }
};

template <std::size_t N> struct fairing_pia_surface : lspia_surface<N>
{
    using vec_t = geom::vec<N, double>; // P
    using lspia_surface<N>::uDegree;
    using lspia_surface<N>::vDegree;
    using lspia_surface<N>::uPoleNum;
    using lspia_surface<N>::vPoleNum;
    using lspia_surface<N>::data;
    using lspia_surface<N>::params;
    using lspia_surface<N>::surface;

    std::vector<double> wStretch; // 每个控制点的拉伸权重
    std::vector<double> wBending; // 每个控制点的光顺权重

  protected:
    // 迭代格式 P' = A * P + B
    Eigen::SparseMatrix<double> m_A;
    Eigen::MatrixXd m_B;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        lspia_surface<N>::init();

        auto B0 = coeff_matrix::surface_least_square(params, surface->knots(0), surface->knots(1), uDegree, vDegree);
        auto [B1, B2] = coeff_matrix::surface_stretch_bending(surface->knots(0), surface->knots(1), uDegree, vDegree);
        auto [A, B] = fairing_pia_curve<N>::coeff_matrices(B0, B1, B2, wStretch, wBending, data);

        m_A = std::move(A);
        m_B = std::move(B);
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<geom::bspline_surface<N>>
     */
    std::shared_ptr<geom::bspline_surface<N>> next() override
    {
        auto n1 = uPoleNum - 1;
        auto n2 = vPoleNum - 1;
        auto n = (n1 + 1) * (n2 + 1);

        // 填充控制点矩阵
        Eigen::MatrixXd P(n, N);
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto idx = i * (n2 + 1) + j;
                for (std::size_t k = 0; k < N; k++)
                    P(idx, k) = surface->pole(i, j)[k];
            }
        }

        // 计算增量
        P = m_A * P + m_B;
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto idx = i * (n2 + 1) + j;
                vec_t p;
                for (std::size_t k = 0; k < N; k++)
                    p[k] = P(idx, k);
                surface->set_pole(p, i, j);
            }
        }

        return surface;
    }
};

} // namespace alg

} // namespace xy
